import os
import sys
from huggingface_hub import snapshot_download
from huggingface_hub.constants import HF_HUB_CACHE

def download_chinese_model():
    """
    从国内镜像下载BAAI/bge-large-zh-v1.5模型
    """
    print("=== BAAI/bge-large-zh-v1.5 模型下载工具 ===")
    print("该模型是专为中文文本设计的高性能嵌入模型")
    print("模型大小约为1.5GB，请确保有足够的磁盘空间\n")
    
    # 设置环境变量使用国内镜像
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    os.environ['HF_HUB_CACHE'] = os.path.join(os.getcwd(), 'models', 'cache')
    
    print(f"缓存目录: {os.environ['HF_HUB_CACHE']}")
    print("下载目录: ./models/bge-large-zh")
    
    # 确保缓存目录存在
    os.makedirs(os.environ['HF_HUB_CACHE'], exist_ok=True)
    
    print("\n开始下载 BAAI/bge-large-zh-v1.5 模型...")
    print("这可能需要一些时间，请耐心等待...")
    print("(如果下载失败，可以尝试使用 https://hf-mirror.com 网站手动下载)\n")
    
    try:
        # 下载模型到本地目录
        model_path = snapshot_download(
            repo_id="BAAI/bge-large-zh-v1.5",
            local_dir="./models/bge-large-zh",
            resume_download=True,  # 支持断点续传
            local_dir_use_symlinks=False  # 直接复制文件而不是创建符号链接
        )
        print(f"✅ 模型已成功下载到: {model_path}")
        print("\n下一步操作:")
        print("1. 项目会自动使用该模型进行文档向量化")
        print("2. 您可以运行项目测试模型效果")
        return model_path
    except Exception as e:
        print(f"❌ 下载过程中出现错误: {e}")
        print("\n🔧 解决方案:")
        print("1. 手动下载方式:")
        print("   访问 https://hf-mirror.com/BAAI/bge-large-zh-v1.5")
        print("   点击 'Files and versions' 标签")
        print("   下载所有文件到 ./models/bge-large-zh 目录")
        print("\n2. 使用Git LFS方式:")
        print("   git lfs install")
        print("   git clone https://hf-mirror.com/BAAI/bge-large-zh-v1.5 models/bge-large-zh")
        print("\n3. 使用huggingface-cli方式:")
        print("   huggingface-cli download --resume-download BAAI/bge-large-zh-v1.5 --local-dir ./models/bge-large-zh --local-dir-use-symlinks False")
        return None

def check_model_exists():
    """
    检查模型是否已经存在
    """
    model_path = "./models/bge-large-zh"
    if os.path.exists(model_path) and os.path.isdir(model_path):
        required_files = ["config.json", "modules.json", "pytorch_model.bin", "sentence_bert_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"]
        found_files = [f for f in required_files if os.path.exists(os.path.join(model_path, f))]
        
        if len(found_files) > len(required_files) * 0.8:  # 如果找到了80%以上的必要文件
            print(f"✅ 检测到模型已存在于: {model_path}")
            print(f"✅ 检测到 {len(found_files)}/{len(required_files)} 个必要文件")
            return True
        else:
            print(f"⚠️  检测到模型目录 {model_path}，但文件不完整")
            return False
    else:
        print(f"❌ 未检测到模型: {model_path}")
        return False

if __name__ == "__main__":
    print("检查模型是否存在...")
    if check_model_exists():
        print("模型已存在，无需重新下载")
        sys.exit(0)
    
    print("开始下载模型...")
    download_chinese_model()